package com.shanzhi.dify.service.impl;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.TypeReference;
import com.shanzhi.connection.common.ConnectionUtils;
import com.shanzhi.connection.common.HttpRequestType;
import com.shanzhi.dify.entity.*;

import com.shanzhi.dify.service.DatasetService;
import com.shanzhi.dify.service.DifyBaseService;
import com.shanzhi.dify.utils.ParseDataUtils;
import com.shanzhi.dify.utils.RebuildParamsUtils;
import kong.unirest.HttpResponse;
import kong.unirest.Unirest;
import lombok.extern.slf4j.Slf4j;
import okhttp3.*;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import com.shanzhi.dify.common.DifyUrlEnums;
import com.shanzhi.dify.entity.DataSetEntity;
import org.springframework.util.StringUtils;

import java.io.File;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 知识库相关的接口
 */
@Service
@Slf4j
public class DatasetServiceImpl implements DatasetService {

    @Autowired
    private DifyBaseService difyBaseService;


    /**
     * 查看知识库的文档列表
     */
    @Override
    public String getDatasetDocuments(DocumentRequestParam documentRequestParam) {
        Map<String, Object> params = new HashMap<>();
        params.put("page", documentRequestParam.getPage());
        params.put("limit", documentRequestParam.getLimit());
        // 如果参数不为空，那么注入入参
        if (StringUtils.hasText(documentRequestParam.getKeyword())) {
            params.put("keyword", documentRequestParam.getKeyword());
        }
        return difyBaseService.sendDifyRequest(null, null, DifyUrlEnums.DATASETS_DOCUMENTS_URL, HttpRequestType.GET, params, documentRequestParam.getDatasetId());
    }

    /**
     * 通过文件创建文档 - Java 1.8兼容版本
     *
     * @param datasetId 知识库ID
     * @param file      文件路径
     * @return 上传结果
     */
    @Override
    public String uploadFileToDataset(String datasetId, File file) {
        try {
            OkHttpClient client = new OkHttpClient().newBuilder().build();

            // 构建data参数
            Map<String, Object> dataParams = new HashMap<>();
            dataParams.put("indexing_technique", "high_quality");
            dataParams.put("doc_form", "text_model");
            dataParams.put("doc_language", "Chinese Simplified");

            // 设置处理规则为自定义模式
            Map<String, Object> processRule = new HashMap<>();
            Map<String, Object> rules = new HashMap<>();

            // 预处理规则 - Java 1.8兼容写法
            List<Map<String, Object>> preProcessingRules = new ArrayList<Map<String, Object>>() {{
                add(new HashMap<String, Object>() {{
                    put("id", "remove_extra_spaces");
                    put("enabled", true);
                }});
                add(new HashMap<String, Object>() {{
                    put("id", "remove_urls_emails");
                    put("enabled", false);
                }});
            }};

            rules.put("pre_processing_rules", preProcessingRules);

            // 分段规则
            Map<String, Object> segmentation = new HashMap<>();
            segmentation.put("separator", "\n\n\n\n");
            segmentation.put("max_tokens", 1024);
            segmentation.put("chunk_overlap", 50);
            rules.put("segmentation", segmentation);

            processRule.put("rules", rules);
            processRule.put("mode", "custom");
            dataParams.put("process_rule", processRule);

            // 设置检索模型参数
            Map<String, Object> retrievalModel = new HashMap<>();
            retrievalModel.put("search_method", "hybrid_search");
            retrievalModel.put("reranking_enable", true);
            retrievalModel.put("reranking_mode", "reranking_model");

            // Reranking 模型设置
            Map<String, Object> rerankingModel = new HashMap<>();
            rerankingModel.put("reranking_provider_name", "langgenius/tongyi/tongyi");
            rerankingModel.put("reranking_model_name", "gte-rerank");
            retrievalModel.put("reranking_model", rerankingModel);

            // 权重设置
            Map<String, Object> weights = new HashMap<>();
            weights.put("weight_type", "customized");

            Map<String, Object> vectorSetting = new HashMap<>();
            vectorSetting.put("vector_weight", 0.7);
            vectorSetting.put("embedding_provider_name", "");
            vectorSetting.put("embedding_model_name", "");
            weights.put("vector_setting", vectorSetting);

            Map<String, Object> keywordSetting = new HashMap<>();
            keywordSetting.put("keyword_weight", 0.3);
            weights.put("keyword_setting", keywordSetting);

            retrievalModel.put("weights", weights);
            retrievalModel.put("top_k", 7);
            retrievalModel.put("score_threshold_enabled", true);
            retrievalModel.put("score_threshold", 0.4);
            dataParams.put("retrieval_model", retrievalModel);

            // Embedding 模型设置
            dataParams.put("embedding_model", "text-embedding-v3");
            dataParams.put("embedding_model_provider", "langgenius/tongyi/tongyi");

            // 将data参数转换为JSON字符串
            String dataJson = JSON.toJSONString(dataParams);

            // 构建multipart请求体
            RequestBody body = new MultipartBody.Builder().setType(MultipartBody.FORM)
                    .addFormDataPart("data", dataJson)
                    .addFormDataPart("file", file.getName(),
                            RequestBody.create(MediaType.parse("application/octet-stream"), file))
                    .build();

            // 构建请求URL
            String url = difyBaseService.buildUrl(DifyUrlEnums.CREATE_DOCUMENT_FILE_URL, datasetId);

            // 获取API Key
            String apiKey = difyBaseService.resolveApiKey(null);

            // 构建请求
            Request request = new Request.Builder()
                    .url(url)
                    .method("POST", body)
                    .addHeader("Authorization", "Bearer " + apiKey)
                    .addHeader("User-Agent", "DifyJavaClient/1.0")
                    .build();

            // 发送请求并返回结果
            Response response = client.newCall(request).execute();
            return response.body().string();

        } catch (Exception e) {
            log.error("上传文件到知识库失败", e);
            throw new RuntimeException("文件上传失败: " + e.getMessage(), e);
        }
    }


    /**
     * 删除文档
     * @param datasetId
     * @param documentId
     */
    @Override
    public void deleteDocument(String datasetId, String documentId) {
        String url = difyBaseService.buildUrl(DifyUrlEnums.DELETE_DOCUMENT_URL, datasetId, documentId);
        String apiKey = difyBaseService.resolveApiKey(null);
        ConnectionUtils
                .sendInternal(null, url, headers -> headers.put("Authorization", "Bearer " + apiKey), HttpRequestType.DELETE, null, null);
    }




    /**
     * 知识库搜索
     *
     * @return
     */
    @Override
    public String datasetSearch(DataSetEntity dataSetEntity) {
        // 参数创建
        RetrievalModelRequestEntity fixedRetrievalRequest = RebuildParamsUtils.createFixedRetrievalRequest(dataSetEntity.getQuery());
        // 使用知识库APiKey发送请求
        return difyBaseService.sendDifyRequest(fixedRetrievalRequest, null, DifyUrlEnums.DATAS_SEARCH_URL, HttpRequestType.POST, null, dataSetEntity.getDatasetId());
    }

    @Override
    public List<String> parseDatasetSearchResult(DataSetEntity dataSetEntity) {
        String datasetSearch = this.datasetSearch(dataSetEntity);
        // 提取segment content
        return ParseDataUtils.extractSegmentContent(datasetSearch);
    }


    @Override
    public List<DocumentFileEntity> extractSegmentContentForDocument(DataSetEntity dataSetEntity) {
        String datasetSearch = this.datasetSearch(dataSetEntity);
        // 提取segment content
        return ParseDataUtils.extractSegmentContentForDocument(datasetSearch);
    }


    /**
     * 创建知识库--使用该方式来进行创建知识库
     *
     * @return
     */

    @Override
    public String createDataset(DifyDatasetEntity difyDatasetEntity) {
        Map<String, String> params = new HashMap<>();
        params.put("name", difyDatasetEntity.getName());
        params.put("permission", "all_team_members");
        String json = difyBaseService.sendDifyRequest(params, null, DifyUrlEnums.DATASETS_URL, HttpRequestType.POST, null);
        return ParseDataUtils.getDatasetId(json);
    }


    /**
     * 更新知识库信息
     * 用户只能修改知识库名称，其他参数保持不变
     *
     * @param datasetId 知识库ID
     * @param name 新的知识库名称
     * @return 更新结果
     */
    @Override
    public String updateDataset(String datasetId, String name) {
        String url = difyBaseService.buildUrl(DifyUrlEnums.DATASETS_DETAIL_URL, datasetId);
        String apiKey = difyBaseService.resolveApiKey(null);


        Map<String, Object> params = new HashMap<>();
        params.put("name", name);
        HttpResponse<String> response = Unirest.patch(url)
                .header("Authorization", "Bearer " + apiKey)
                .header("User-Agent", "Apifox/1.0.0 (https://apifox.com)")
                .header("Content-Type", "application/json")
                .header("Accept", "*/*")
                .header("Host", "162.14.113.149:8097")
                .header("Connection", "keep-alive")
                .body(JSONObject.toJSONString(params))
                .asString();
        return ParseDataUtils.getDatasetId(response.getBody());
    }

    /**
     * 删除知识库
     *
     * @param datasetId 知识库ID
     * @return 删除结果
     */
    @Override
    public void deleteDataset(String datasetId) {
        String url = difyBaseService.buildUrl(DifyUrlEnums.DATASETS_DETAIL_URL, datasetId);
        String apiKey = difyBaseService.resolveApiKey(null);

        ConnectionUtils
                .sendInternal(null, url, headers -> headers.put("Authorization", "Bearer " + apiKey), HttpRequestType.DELETE, null, null);
    }


    /**
     * 创建知识库-----后面弃用
     *
     * @return
     */

    @Override
    public String createDataset(DifyKnowledgeEntity knowledgeEntity) {
        String json = difyBaseService.sendDifyRequest(knowledgeEntity, null, DifyUrlEnums.DATASETS_URL, HttpRequestType.POST, null);
        return ParseDataUtils.getDatasetId(json);
    }

    /**
     * 上传文件获取上传文件 /datasets/{dataset_id}/documents/{document_id}/upload-file
     *
     * @return
     */

    @Override
    public Map<String, Object> getUploadFile(String datasetId, List<String> documentIds) {
        Map<String, Object> map = new HashMap<>();
        for (String documentId : documentIds) {
            String json = difyBaseService.sendDifyRequest(null, null, DifyUrlEnums.UPLOAD_DOCUMENTS_URL, HttpRequestType.GET, null, datasetId, documentId);
            map.put(documentId, JSON.parseObject(json, new TypeReference<Map<String, Object>>() {
            }));
        }
        return map;
    }

    /**
     * 获取查询文档子分段(块)
     * /datasets/{dataset_id}/documents/{document_id}/segments/{segment_id}/child_chunks
     */
    @Override
    public Map<String, Object> getSegmentChildChunks(String datasetId, String documentId, List<String> segmentIds) {
        Map<String, Object> map = new HashMap<>();
        for (String segmentId : segmentIds) {
            String json = difyBaseService.sendDifyRequest(null, null, DifyUrlEnums.DOCUMENTS_SEGMENTS_CHILD_CHUNKS_URL, HttpRequestType.GET, null, datasetId, documentId, segmentId);
            map.put(documentId, JSON.parseObject(json, new TypeReference<Map<String, Object>>() {
            }));
        }
        return map;
    }


}
